﻿#ifndef __ROAD_LINE_GRAY_MODEL_H__
#define __ROAD_LINE_GRAY_MODEL_H__

#include "JsColorModel.h"
#include "base64.h"
#include "models/SoftMaxModel.h"

using SoftMaxModel = darknet::SoftMaxModel;
class RoadLineGrayModel : public SoftMaxModel
{
    using Options   = darknet::Options;
    using DataImage = darknet::DataImage;

public:
    static RoadLineGrayModel *GetInstance()
    {
        static RoadLineGrayModel s_instance;
        return &s_instance;
    }

public:
    RoadLineGrayModel()
    {
        m_max_epochs      = 4000;
        m_batch           = 4;
        m_auto_stop_train = false;
        m_jitter          = 0.f;
        m_distort         = false;
        m_augment         = false;
        m_model_name      = "roadlines_gray";

#ifdef __EMSCRIPTEN__
        std::string model_filepath  = "/models/" + GetDefaultModelFileName();
        std::string model_directory = "/darknet/models";
#else
        std::string model_filepath  = "D:/workspace/cpp/darknet/data/roadlines/models/" + GetDefaultModelFileName();
        std::string model_directory = "D:/workspace/cpp/darknet/data/roadlines/models";
#endif
        if (std::filesystem::exists(model_filepath) && LoadModelFromFile(model_filepath))
        {
            std::cout << "load prebuilt model from " << model_filepath << std::endl;
        }
        else
        {
            LoadNet();
        }
        SetModelDirectory(model_directory);

        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行左转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行右转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行左转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行右转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转45度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转90度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转45度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转90度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(禁止通行)"));
    }

    virtual void LoadNet()
    {
        if (GetNet()->IsLoaded()) return;

        auto net     = GetNet();
        auto options = std::make_shared<std::vector<std::shared_ptr<Options>>>();
        auto option  = std::make_shared<Options>("net");
        option->SetOption("width", 64);
        option->SetOption("height", 36);
        option->SetOption("channels", 1);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 32);
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", 1);
        option->SetOption("activation", "linear");
        options->push_back(option);

        option = std::make_shared<Options>("softmax");
        options->push_back(option);

        net->Load(options);
        net->SetBatch(GetBatch());
    }

    virtual void LoadDataSet()
    {
        SoftMaxModel::LoadDataSet();

#ifndef __EMSCRIPTEN__
        auto auto_directory = GetDataSetDirectory() + "/auto";
        std::filesystem::create_directories(auto_directory);
        auto dataset = GetDataSet();
        for (auto it = dataset->begin(); it != dataset->end(); it++)
        {
            std::string image_path = it->first;
            auto img               = std::make_shared<DataImage>();
            if (!img->Load(image_path))
            {
                std::cout << "invalid image path: " << image_path << std::endl;
                continue;
            }

            img = HandleInputImage(img);

            auto path = std::filesystem::path(image_path);
            img->Save((std::filesystem::path(auto_directory) / path.filename().replace_extension(".jpg")).string());
        }
#endif
    }

    virtual std::shared_ptr<DataImage> HandleInputImage(std::shared_ptr<DataImage> input_image)
    {
        auto img             = input_image->Scale(0.4, 0.4); // 320 * 240 => 128 * 96
        auto step            = 2;                            // 128 * 96 => 64 * 48
        auto threshold       = step * step * 2 / 4;          // 判定阈值
        auto output_image    = std::make_shared<DataImage>(GetNetWidth(), GetNetHeight(), GetNetChannels());
        auto output_channels = GetNetChannels();
        auto color_model     = JsColorModel::GetInstance();
        auto roadsign_index  = color_model->GetRedClassifyIndex();
        auto roadline_index  = color_model->GetWhiteClassifyIndex();
        auto roadedge_index  = color_model->GetBlackClassifyIndex();
        auto datas           = img->GetImageRGBData()->data();
        auto width           = img->GetWidth();
        auto height          = img->GetHeight();
        auto offset_y        = height - GetNetHeight() * step;
        for (int y = offset_y; y < height; y += step)
        {
            for (int x = 0; x < width; x += step)
            {
                int roadline_count = 0;
                int roadedge_count = 0;
                for (int i = 0; i < step && roadline_count < threshold && roadedge_count < threshold; i++)
                {
                    for (int j = 0; j < step && roadline_count < threshold && roadedge_count < threshold; j++)
                    {
                        auto index       = ((y + i) * width + (x + j)) * 3;
                        auto color_index = color_model->Predict(datas[index], datas[index + 1], datas[index + 2]);
                        if (roadline_index == color_index || roadsign_index == color_index) roadline_count++;
                        if (roadedge_index == color_index) roadedge_count++;
                    }
                }
                auto output_x = x / step;
                auto output_y = (y - offset_y) / step;
                for (int k = 0; k < output_channels; k++)
                {
                    if (roadline_count >= threshold)
                    {
                        output_image->SetPixel(output_x, output_y, k, 1.0f);
                    }
                    else if (roadedge_count >= threshold)
                    {
                        output_image->SetPixel(output_x, output_y, k, 0.5f);
                    }
                    else
                    {
                        output_image->SetPixel(output_x, output_y, k, 0.0f);
                    }
                }
            }
        }
        return output_image;
    }

public:
    std::string FixColorData(int *colors, int width, int height)
    {
        int white_color = 1;
        int black_color = 0;
        int other_color = -1;

        auto x_offset          = 0;
        auto y_offset          = height / 4;
        auto center_x          = x_offset + (width - x_offset) / 2;
        auto center_y          = y_offset + (height - y_offset) / 2;
        auto total_count       = (width - x_offset) * (height - y_offset) / 2;
        auto white_color_count = GetColorCount(colors, white_color, x_offset, y_offset + (height - y_offset) / 2, width - x_offset, (height - y_offset) / 2, width, height);
        auto black_color_count = GetColorCount(colors, black_color, x_offset, y_offset + (height - y_offset) / 2, width - x_offset, (height - y_offset) / 2, width, height);
        auto other_color_count = GetColorCount(colors, other_color, x_offset, y_offset + (height - y_offset) / 2, width - x_offset, (height - y_offset) / 2, width, height);
        std::cout << "white_color_count = " << white_color_count << " black_color_count = " << black_color_count << " other_color_count = " << other_color_count << std::endl;
        if ((black_color_count + other_color) > (total_count * 2 / 3)) return reinterpret_cast<const char*>(u8R"(禁止通行)");
        if (white_color_count < (total_count * 1 / 3)) return reinterpret_cast<const char*>(u8R"(禁止通行)");

        auto [white_color_center_x, white_color_center_y] = GetColorCenter(colors, white_color, x_offset, y_offset, width - x_offset, height - y_offset, width, height);
        if (white_color_center_x == -1 && white_color_center_y == -1) return reinterpret_cast<const char*>(u8R"(禁止通行)");
        auto black_color_center_x = 0;
        auto black_color_center_y = 0;
        if (white_color_center_x < center_x)
        {
            // 路面在左, 路线在右
            auto paris           = GetColorCenter(colors, black_color, white_color_center_x, y_offset, width - white_color_center_x, height - y_offset, width, height);
            black_color_center_x = std::lround(paris.first);
            black_color_center_y = std::lround(paris.second);
        }
        else
        {
            // 路面在右, 路线在左
            auto paris           = GetColorCenter(colors, black_color, x_offset, y_offset, white_color_center_x - x_offset, height - y_offset, width, height);
            black_color_center_x = std::lround(paris.first);
            black_color_center_y = std::lround(paris.second);
        }
        auto x = black_color_center_x < white_color_center_x ? black_color_center_x : white_color_center_x;
        auto y = y_offset;
        auto w = black_color_center_x < white_color_center_x ? (white_color_center_x - black_color_center_x) : (black_color_center_x - white_color_center_x);
        // auto h = (black_color_center_y < white_color_center_y ? white_color_center_y : black_color_center_y) - y_offset;
        auto h = height - y_offset;
        std::cout << "white_color_center_x = " << white_color_center_x << " white_color_center_y = " << white_color_center_y << std::endl;
        std::cout << "black_color_center_x = " << black_color_center_x << " black_color_center_y = " << black_color_center_y << std::endl;
        std::cout << "x = " << x << " y = " << y << " w = " << w << " h = " << h << std::endl;
        auto [p1_x, p1_y] = GetColorCenter(colors, black_color, x, y, w / 2, h, width, height);
        auto [p2_x, p2_y] = GetColorCenter(colors, black_color, x + w / 2, y, w / 2, h, width, height);
        std::cout << "p1_x = " << p1_x << " p1_y = " << p1_y << " p2_x = " << p2_x << " p2_y = " << p2_y << std::endl;
        if (std::abs(p1_y - p2_y) <= 1)
        {
            if (p1_y > (height * 3 / 4)) return reinterpret_cast<const char*>(u8R"(禁止通行)");
            return black_color_center_x < white_color_center_x ? reinterpret_cast<const char*>(u8R"(右转90度)") : reinterpret_cast<const char*>(u8R"(左转90度)");
        }
        else if (p1_x < 0 || p2_x < 0 || std::abs(p1_x - p2_x) <= 1)
        {
            auto offset = white_color_center_x - center_x;
            if (offset >= 2)
            {
                return reinterpret_cast<const char*>(u8R"(直行偏右)");
            }
            else if (offset <= -2)
            {
                return reinterpret_cast<const char*>(u8R"(直行偏左)");
            }
            else
            {
                return reinterpret_cast<const char*>(u8R"(直行)");
            }
        }
        else
        {
            p1_y        = height - p1_y - 1;
            p2_y        = height - p2_y - 1;
            auto slope  = (p1_y - p2_y) * 1.0f / (p1_x - p2_x);
            int degree  = std::atan(slope) * 180 / M_PI;
            auto offset = std::abs(white_color_center_x - center_x);
            if (offset > 5 && std::abs(degree) > 45) degree = degree < 0 ? -45 : 45;
            return (degree < 0 ? reinterpret_cast<const char*>(u8R"(左转)") : reinterpret_cast<const char*>(u8R"(右转)")) + std::to_string(90 - std::abs(degree)) + reinterpret_cast<const char*>(u8R"(度)");
        }
    }

    int GetColorCount(int *colors, int color, int x, int y, int width, int height, int max_width, int max_height)
    {
        int count = 0;
        for (int i = 0; i < width; i++)
        {
            for (int j = 0; j < height; j++)
            {
                int ix = i + x;
                int jy = j + y;
                if (ix < 0 || ix >= max_width || jy < 0 || jy >= max_height) continue;
                int index = jy * max_width + ix;
                if (colors[index] == color)
                {
                    count++;
                }
            }
        }
        return count;
    }

    std::pair<float, float> GetColorCenter(int *colors, int color, int x, int y, int width, int height, int max_width, int max_height)
    {
        int avg_x    = 0;
        int avg_y    = 0;
        int avg_size = 0;
        for (int i = 0; i < width; i++)
        {
            for (int j = 0; j < height; j++)
            {
                int ix = i + x;
                int jy = j + y;
                if (ix < 0 || ix >= max_width || jy < 0 || jy >= max_height) continue;
                int index = jy * max_width + ix;
                if (colors[index] == color)
                {
                    avg_x += ix;
                    avg_y += jy;
                    avg_size++;
                }
            }
        }
        if (avg_size > 0)
        {
            return std::make_pair(avg_x * 1.0f / avg_size, avg_y * 1.0f / avg_size);
        }
        else
        {
            return std::make_pair(-1.f, -1.f);
        }
    }

protected:
    static std::string LoadFile(const std::string &path)
    {
        std::ifstream ifs(path, std::ios::binary);
        if (ifs)
        {
            std::string text((std::istreambuf_iterator<char>(ifs)), (std::istreambuf_iterator<char>()));
            ifs.close();
            return text;
        }
        else
        {
            std::cout << "invalid path: " << path << std::endl;
            return "";
        }
    }
};

#endif